import pandas as pd
from mlxtend.frequent_patterns import apriori, association_rules
from mlxtend.preprocessing import TransactionEncoder
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from tqdm import tqdm  # 用于添加进度条
import time

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False


def load_and_preprocess_data(file_path):
    """加载并预处理数据，构建用户-商品矩阵"""
    try:
        print("开始加载数据...")
        # 读取Excel文件
        df = pd.read_excel(file_path)
        print(f"数据加载成功，共{len(df)}条交易记录")
        print(f"数据包含以下列：{list(df.columns)}")

        # 确保关键列存在
        required_columns = ['customer_id', 'product_id']
        for col in required_columns:
            if col not in df.columns:
                raise ValueError(f"数据中缺少必要的列 '{col}'，分析无法进行")

        # 构建用户-商品矩阵 (稀疏矩阵形式)
        print("正在构建用户-商品矩阵...")
        unique_customers = df['customer_id'].unique()
        unique_products = df['product_id'].unique()

        # 使用tqdm显示进度
        user_item_matrix = pd.DataFrame(index=unique_customers, columns=unique_products, data=0)

        # 填充矩阵 - 使用进度条
        print("正在填充用户-商品矩阵...")
        for customer_id in tqdm(unique_customers, desc="填充矩阵"):
            customer_products = df[df['customer_id'] == customer_id]['product_id'].unique()
            user_item_matrix.loc[customer_id, customer_products] = 1

        print(f"用户-商品矩阵构建完成，包含{user_item_matrix.shape[0]}个用户和{user_item_matrix.shape[1]}个商品")

        return user_item_matrix, df
    except Exception as e:
        print(f"数据加载失败: {e}")
        raise


def analyze_user_item_relations(user_item_matrix, min_support=0.01, min_confidence=0.5):
    """分析用户与商品之间的关联关系"""
    # 转置矩阵，使行为交易(用户)，列为商品
    matrix_t = user_item_matrix.T

    # 使用Apriori算法挖掘频繁项集
    print("正在挖掘频繁项集...")
    frequent_itemsets = apriori(matrix_t, min_support=min_support, use_colnames=True)

    # 生成关联规则
    print("正在生成关联规则...")
    rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=min_confidence)

    print(f"发现{len(frequent_itemsets)}个频繁项集和{len(rules)}条关联规则")
    return frequent_itemsets, rules


def filter_rules_for_user(rules, target_user, user_item_matrix):
    """筛选与目标用户相关的关联规则"""
    # 将用户ID转换为字符串格式（与规则中的格式一致）
    target_user_str = str(target_user)

    # 筛选前件或后件中包含目标用户的规则
    print(f"正在筛选与用户 {target_user} 相关的关联规则...")
    user_rules = rules[rules['antecedents'].apply(lambda x: target_user_str in x) |
                       rules['consequents'].apply(lambda x: target_user_str in x)]

    print(f"找到{len(user_rules)}条与用户 {target_user} 相关的关联规则")

    # 提取关联的商品信息
    user_product_relations = []

    # 使用tqdm显示处理进度
    print("正在分析关联规则中的用户-商品关系...")
    for _, rule in tqdm(user_rules.iterrows(), total=len(user_rules), desc="分析规则"):
        antecedents = set(rule['antecedents'])
        consequents = set(rule['consequents'])

        # 确定规则中哪些是用户，哪些是商品
        antecedent_users = {u for u in antecedents if u in user_item_matrix.index.astype(str)}
        antecedent_products = antecedents - antecedent_users

        consequent_users = {u for u in consequents if u in user_item_matrix.index.astype(str)}
        consequent_products = consequents - consequent_users

        # 如果规则涉及用户与商品的关系
        if (antecedent_users or consequent_users) and (antecedent_products or consequent_products):
            # 用户购买商品的规则 (用户 → 商品)
            if antecedent_users and consequent_products:
                for user in antecedent_users:
                    for product in consequent_products:
                        user_product_relations.append({
                            'user_id': int(user),
                            'product_id': int(product),
                            'confidence': rule['confidence'],
                            'lift': rule['lift'],
                            'support': rule['support'],
                            'relation_type': '用户→商品'
                        })

            # 商品→用户的规则 (表示购买该商品的用户倾向于也是目标用户)
            if consequent_users and antecedent_products:
                for user in consequent_users:
                    for product in antecedent_products:
                        user_product_relations.append({
                            'user_id': int(user),
                            'product_id': int(product),
                            'confidence': rule['confidence'],
                            'lift': rule['lift'],
                            'support': rule['support'],
                            'relation_type': '商品→用户'
                        })

    return pd.DataFrame(user_product_relations)


def analyze_specific_user_product(user_id, product_id, user_item_matrix, rules_df):
    """分析特定用户与特定商品之间的关联度"""
    # 检查用户是否存在
    if user_id not in user_item_matrix.index:
        print(f"用户 {user_id} 不存在于数据中")
        return None

    # 筛选与该用户相关的规则
    print(f"正在筛选用户 {user_id} 相关的规则...")
    user_rules = rules_df[rules_df['user_id'] == user_id]

    if len(user_rules) == 0:
        print(f"没有找到与用户 {user_id} 相关的关联规则")
        return None

    # 筛选与特定商品相关的规则
    print(f"正在分析用户 {user_id} 与商品 {product_id} 的关联度...")
    user_product_rules = user_rules[user_rules['product_id'] == product_id]

    if len(user_product_rules) > 0:
        print(f"\n用户 {user_id} 与商品 {product_id} 的关联度分析:")
        for _, rule in user_product_rules.iterrows():
            print(f"- 关联类型: {rule['relation_type']}")
            print(f"  支持度: {rule['support']:.4f} - 表示同时涉及用户和商品的交易比例")
            print(f"  置信度: {rule['confidence']:.4f} - 表示用户关联该商品的可能性")
            print(f"  提升度: {rule['lift']:.4f} - 表示用户与商品的关联强度（>1表示正相关）")
        return user_product_rules
    else:
        print(f"没有找到用户 {user_id} 与商品 {product_id} 直接相关的规则")
        print("但可以查看用户最相关的商品:")
        top_products = user_rules.sort_values('confidence', ascending=False).head(5)
        for _, product in top_products.iterrows():
            print(f"- 商品 {product['product_id']}: 置信度={product['confidence']:.4f}, 提升度={product['lift']:.4f}")
        return top_products


def visualize_user_product_network(user_id, user_rules, top_n=10):
    """可视化用户与商品的关联网络"""
    if user_rules is None or len(user_rules) == 0:
        return

    # 按置信度排序，取前n个商品
    print(f"正在准备用户 {user_id} 的关联网络数据...")
    top_products = user_rules.sort_values('confidence', ascending=False).head(top_n)

    if len(top_products) == 0:
        return

    import networkx as nx

    G = nx.DiGraph()

    # 添加节点
    G.add_node(user_id, type='user', label=f'用户{user_id}')

    # 添加商品节点和边 - 使用进度条
    print(f"正在构建用户 {user_id} 的关联网络图...")
    for _, product in tqdm(top_products.iterrows(), total=len(top_products), desc="添加节点"):
        product_id = product['product_id']
        confidence = product['confidence']
        lift = product['lift']

        G.add_node(product_id, type='product', label=f'商品{product_id}')
        G.add_edge(user_id, product_id, weight=confidence,
                   label=f'置信:{confidence:.2f}\n提升:{lift:.2f}')

    # 绘制网络图
    print(f"正在绘制用户 {user_id} 的关联网络图...")
    plt.figure(figsize=(12, 8))

    # 设置节点位置
    pos = nx.spring_layout(G, k=0.3)

    # 绘制节点
    user_nodes = [n for n, d in G.nodes(data=True) if d['type'] == 'user']
    product_nodes = [n for n, d in G.nodes(data=True) if d['type'] == 'product']

    nx.draw_networkx_nodes(G, pos, nodelist=user_nodes, node_size=1000,
                           node_color='lightblue', alpha=0.8)
    nx.draw_networkx_nodes(G, pos, nodelist=product_nodes, node_size=800,
                           node_color='lightgreen', alpha=0.8)

    # 绘制边
    edges = G.edges()
    weights = [G[u][v]['weight'] * 3 for u, v in edges]

    nx.draw_networkx_edges(G, pos, width=weights, alpha=0.6,
                           edge_color='gray', arrowsize=20)

    # 添加标签
    node_labels = {n: d['label'] for n, d in G.nodes(data=True)}
    edge_labels = {(u, v): G[u][v]['label'] for u, v in edges}

    nx.draw_networkx_labels(G, pos, labels=node_labels, font_size=12)
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=10)

    plt.title(f'用户 {user_id} 与商品的关联网络', fontsize=15)
    plt.axis('off')
    plt.tight_layout()
    plt.savefig(f'user_{user_id}_network.png', dpi=300, bbox_inches='tight')
    plt.close()

    print(f"用户 {user_id} 的关联网络已保存为 user_{user_id}_network.png")


def main():
    # 文件路径 - 需要替换为实际文件路径
    file_path = 'A_xxx.xlsx'

    # 加载和预处理数据
    user_item_matrix, df = load_and_preprocess_data(file_path)

    # 分析用户与商品的关联关系
    frequent_itemsets, rules = analyze_user_item_relations(user_item_matrix,
                                                           min_support=0.01,  # 降低支持度阈值以发现更多用户关联
                                                           min_confidence=0.3)  # 降低置信度阈值

    # 转换规则格式，便于分析
    print("正在转换关联规则格式...")
    rules_df = pd.DataFrame()

    # 使用tqdm显示进度
    for _, rule in tqdm(rules.iterrows(), total=len(rules), desc="处理规则"):
        antecedents = list(rule['antecedents'])
        consequents = list(rule['consequents'])

        # 检查前件和后件是否包含用户或商品
        antecedent_users = [int(u) for u in antecedents if u in user_item_matrix.index.astype(str)]
        antecedent_products = [int(p) for p in antecedents if p not in user_item_matrix.index.astype(str)]

        consequent_users = [int(u) for u in consequents if u in user_item_matrix.index.astype(str)]
        consequent_products = [int(p) for p in consequents if p not in user_item_matrix.index.astype(str)]

        # 用户→商品规则
        if antecedent_users and consequent_products:
            for user in antecedent_users:
                for product in consequent_products:
                    rules_df = pd.concat([rules_df, pd.DataFrame({
                        'user_id': [user],
                        'product_id': [product],
                        'support': [rule['support']],
                        'confidence': [rule['confidence']],
                        'lift': [rule['lift']],
                        'relation_type': ['用户→商品']
                    })], ignore_index=True)

        # 商品→用户规则
        if consequent_users and antecedent_products:
            for user in consequent_users:
                for product in antecedent_products:
                    rules_df = pd.concat([rules_df, pd.DataFrame({
                        'user_id': [user],
                        'product_id': [product],
                        'support': [rule['support']],
                        'confidence': [rule['confidence']],
                        'lift': [rule['lift']],
                        'relation_type': ['商品→用户']
                    })], ignore_index=True)

    # 示例：分析特定用户与特定商品的关联度
    target_user_id = 1  # 替换为你感兴趣的用户ID
    target_product_id = 101  # 替换为你感兴趣的商品ID

    user_product_relations = analyze_specific_user_product(
        target_user_id, target_product_id, user_item_matrix, rules_df)

    # 可视化用户与商品的关联网络
    visualize_user_product_network(target_user_id, user_product_relations)

    # 示例：为用户推荐商品
    print("\n基于关联规则的商品推荐:")
    recommend_products_for_user(target_user_id, user_item_matrix, rules_df)


def recommend_products_for_user(user_id, user_item_matrix, rules_df, top_n=5):
    """基于关联规则为用户推荐商品"""
    # 检查用户是否存在
    if user_id not in user_item_matrix.index:
        print(f"用户 {user_id} 不存在于数据中")
        return

    # 获取用户已购买的商品
    user_purchased = set(user_item_matrix.columns[user_item_matrix.loc[user_id] > 0])

    # 筛选与该用户相关的规则
    print(f"正在为用户 {user_id} 筛选推荐商品...")
    user_rules = rules_df[rules_df['user_id'] == user_id]

    if len(user_rules) == 0:
        print(f"没有找到与用户 {user_id} 相关的关联规则，无法进行推荐")
        return

    # 按置信度和提升度排序
    sorted_rules = user_rules.sort_values(['confidence', 'lift'], ascending=False)

    # 收集推荐商品
    recommended_products = []

    # 使用tqdm显示进度
    for _, rule in tqdm(sorted_rules.iterrows(), total=len(sorted_rules), desc="生成推荐"):
        product_id = rule['product_id']

        # 如果商品未被用户购买过，则推荐
        if product_id not in user_purchased:
            recommended_products.append({
                'product_id': product_id,
                'confidence': rule['confidence'],
                'lift': rule['lift']
            })

            # 达到推荐数量上限则停止
            if len(recommended_products) >= top_n:
                break

    if recommended_products:
        print(f"为用户 {user_id} 推荐的商品:")
        for i, product in enumerate(recommended_products, 1):
            print(
                f"{i}. 商品ID: {product['product_id']}, 置信度: {product['confidence']:.4f}, 提升度: {product['lift']:.4f}")
    else:
        print(f"没有找到用户 {user_id} 未购买过的关联商品")


if __name__ == "__main__":
    main()